package com.btl.service.ai.support.rag.embedding.strategy.ingestor;

import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingModelTypeEnum;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingStoreIngestorStrategy;
import com.btl.service.ai.support.rag.store.vector.collections.VectorStoreQdrantFactory;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.IngestionResult;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

import java.util.List;

@Slf4j
@Service
public class BgeSmallZhV15EmbeddingStoreIngestorStrategy implements EmbeddingStoreIngestorStrategy {

    private final EmbeddingModel begSmallZhV15Model;
    private final VectorStoreQdrantFactory vectorStoreQdrantFactory;

    public BgeSmallZhV15EmbeddingStoreIngestorStrategy(@Qualifier("begSmallZhV15Model") EmbeddingModel begSmallZhV15Model, VectorStoreQdrantFactory vectorStoreQdrantFactory) {
        this.begSmallZhV15Model = begSmallZhV15Model;
        this.vectorStoreQdrantFactory = vectorStoreQdrantFactory;
    }

    @Override
    public EmbeddingModel embeddingModel() {
        return begSmallZhV15Model;
    }

    @Override
    public EmbeddingStore<TextSegment> embeddingStore(String collectionName) {
        return vectorStoreQdrantFactory.getQdrantClient(collectionName.concat("_")
                .concat(String.valueOf(begSmallZhV15Model.dimension())));
    }

    @Override
    public IngestionResult embed(String collectionName, DocumentSplitter documentSplitter, List<Document> documents) {
        log.info("开始导入Vector集合:{}-记录数:{}", collectionName, documents.size());
        EmbeddingStore<TextSegment> embeddingStore = this.embeddingStore(collectionName);
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .embeddingModel(begSmallZhV15Model)
                .documentSplitter(null)
                .embeddingStore(embeddingStore)
                .build();
        List<String> ids = documents.stream()
                .map(document -> {
                    return document.metadata()
                            .getString("vectorId");
                })
                .toList();
        IngestionResult ingest = ingestor.ingest(ids,documents);
        log.info("完成导入Vector集合,ingest: {}", ingest.tokenUsage());
        return ingest;
    }

    @Override
    public IngestionResult embed(String collectionName, List<Document> documents) {
        return this.embed(collectionName, null, documents);
    }

    @Override
    public EmbeddingModelTypeEnum modelType() {
        return EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL;
    }
}
